import csv
import os


def load_entity_data(ori_file, video_root):
    """
    从 CSV 文件生成 entity_data, speech_data, ts_to_entity, entity_list
    不依赖缓存文件。
    """
    entity_data = {}
    speech_data = {}
    ts_to_entity = {}

    def csv_to_list(csv_path):
        with open(csv_path, 'r', encoding='utf-8') as f:
            reader = csv.reader(f)
            return list(reader)

    def postprocess_speech_label(speech_label):
        speech_label = int(speech_label)
        if speech_label == 2:  # 2 = SPEAKING_NOT_AUDIBLE
            speech_label = 0
        return speech_label

    def cache_entity_data(csv_file_path):
        entity_set = set()
        csv_data = csv_to_list(csv_file_path)
        csv_data.pop(0)  # 去除表头
        for csv_row in csv_data:
            video_id = csv_row[0]
            entity_id = csv_row[-3]
            timestamp = csv_row[1]
            speech_label = postprocess_speech_label(csv_row[-2])

            # 更新 entity_data
            entity_data.setdefault(video_id, {}).setdefault(entity_id, {})[timestamp] = speech_label
            entity_set.add((video_id, entity_id))

            # 更新 speech_data
            speech_data.setdefault(video_id, {}).setdefault(timestamp, 0)
            speech_data[video_id][timestamp] = max(speech_data[video_id][timestamp], speech_label)

        return entity_set

    def entity_list_postprocessing(entity_set, video_root):
        print('初始化实体列表，总实体数：', len(entity_set))

        # 过滤磁盘上不存在的实体
        for video_id, entity_id in entity_set.copy():
            exist_entity = os.path.join(video_root, video_id, entity_id)
            if not os.path.exists(exist_entity):
                entity_set.remove((video_id, entity_id))
        print('过滤未下载实体后，总实体数：', len(entity_set))

        # 过滤图片不完整的实体
        for video_id, entity_id in entity_set.copy():
            dir_path = os.path.join(video_root, video_id, entity_id)
            if len(os.listdir(dir_path)) != len(entity_data[video_id][entity_id]):
                entity_set.remove((video_id, entity_id))
        print('过滤图片不完整实体后，总实体数：', len(entity_set))

        # 转换为排序列表
        entity_list = sorted(list(entity_set))

        # 构建 ts_to_entity
        for video_id, entity_id in entity_set:
            for timestamp in entity_data[video_id][entity_id]:
                ts_to_entity.setdefault(video_id, {}).setdefault(timestamp, []).append(entity_id)

        return entity_list

    def clean_entity_data(entity_list, entity_data):
        valid_entities = set(entity_list)
        for video_id in list(entity_data.keys()):
            for entity_id in list(entity_data[video_id].keys()):
                if (video_id, entity_id) not in valid_entities:
                    del entity_data[video_id][entity_id]
            if not entity_data[video_id]:
                del entity_data[video_id]
        return entity_data

    # 主流程
    entity_set = cache_entity_data(ori_file)
    entity_list = entity_list_postprocessing(entity_set, video_root)
    entity_data = clean_entity_data(entity_list, entity_data)

    return entity_data, ts_to_entity, speech_data, entity_list